import gym
import numpy as np

from pyglet_render.viewer import Viewer
from ship_envs.core.position import get_random_position, get_random_target_range_position
from ship_envs.core.ship import Ship
from ship_envs.core.ship_action_continuous import ShipActionContinus
from ship_envs.core.world import World
from ship_envs.core.model_transformer import model2geom
from ship_envs.core import action
from ship_envs.core import observation


class SingleShipContinuousEnv(gym.Env):
    def __init__(self):
        super().__init__()
        self.world = World()
        self.world.add_ship(ShipActionContinus())
        # 船与目标点之间的距离
        self.distance2targets = [0.2, 0.4, 0.8]
        # 每个距离产生的概率
        self.distance_probs = [0.5, 0.2, 0.3]

        # 执行步数限制为500
        self.step_count = 0
        self.step_count_max = 300
        # spaces
        # 状态空间 此处定义为一个向量和一个坐标点
        # 向量：船->目标点  的向量，x,y范围均为[-1,1]
        # 坐标：船的中心点左边，x,y范围均为[0,1]
        # self.observation_space = gym.spaces.Box(low=np.array([-1, -1, 0, 0]), high=np.array([1, 1, 1, 1]), shape=(4,),
        #                                         dtype=np.float32)
        # 以图像为状态空间
        self.observation_space = observation.get_space()

        # 动作空间 此处定义为 动力加速度和转向加速度
        # 动力加速度：控制前进和减速
        # 转向加速度：控制转向
        # 此处使用DQN测试环境，设置动作空间为离散，仅控制转向
        # 0: 不转向 1：逆时针5° 2：顺时针5°
        self.action_space = action.get_continuous_space()

        # 渲染视图
        self.viewer = None

        self.reward_threshold = 0.8

        self.show_fairway = False

    def get_ship(self):
        # 单智能体场景
        return self.world.get_agent_ship(0)

    def set_show_fairway(self, show):
        self.show_fairway = show

    def step(self, action):
        self.step_count += 1
        self.world.step([action])
        obs = self.observation_callback(self.world, self.get_ship())
        reward = self.reward_callback(self.world, self.get_ship())
        done = self.done_callback(self.world, self.get_ship())
        info = self.info_callback(self.world, self.get_ship())
        return obs, reward, done, info

    def reset(self):
        # 重置船位置及目标点位置
        ship = self.get_ship()
        ship.cur_speed = 0.004
        ship.cur_rotate_speed = 0
        pos = get_random_position()
        ship.set_pos(pos)
        ship.target.set_pos(get_random_target_range_position(pos, self.distance2targets, self.distance_probs))

        # 重置步数
        self.step_count = 0
        # 刷新obs
        self.world.refresh()
        # 单智能体，obs为整个world的obs  全观测
        return self.observation_callback(self.world, ship)

    def render(self, mode="human"):
        if not self.viewer:
            self.viewer = Viewer(700, 700)
            self.viewer.set_bounds(0, 1, 0, 1)
        self.viewer.geoms = []
        entitys = self.world.get_entitys()
        for entity in entitys:
            self.viewer.add_geom(model2geom(entity.model))
            if isinstance(entity, Ship):
                self.viewer.add_geom(model2geom(entity.target))
                # 显示航道
                if self.show_fairway and entity.fairway:
                    self.viewer.add_geom(model2geom(entity.fairway.get_model()))
        self.viewer.render()

    def seed(self, seed=None):
        np.random.seed(seed)

    def steps_out_range(self):
        return self.step_count > self.step_count_max

    def reward_callback(self, world, ship):
        if ship.arrive_target():
            return 100
        if world.has_collision() or ship.out_range():
            return -100
        if self.steps_out_range():
            return -100
        reward = 0
        if ship.cur_rotate_speed > 0:
            reward += -1
        elif ship.cur_rotate_speed < 0:
            reward += -2
        if ship.from_away2close():
            reward += 1
        elif ship.from_close2away():
            reward -= 2
        return reward

    def observation_callback(self, world, ship):
        return world.get_obs()

    def done_callback(self, world, ship):
        return world.has_collision() or ship.arrive_target() or ship.out_range() or self.steps_out_range()

    def info_callback(self, world, ship):
        return {}